#include "kHashmap.h"

#include "sysapi.h"
#include "kvector.h"
#include "klist.h"

#include <memory.h>
#include <stdio.h>

#define LOAD_FACTOR                     1
#define INITIAL_BUCKETS_AMOUNT          8

#define BUCKETS_MULTIPLIER_GROWING_EDGE 512
#define FAST_GROWING_BUCKETS_MULTIPLIER 8
#define SLOW_GROWING_BUCKETS_MULTIPLIER 2

typedef unsigned char uint8_t;

typedef struct DataHeader {
    void*              key;
    size_t             keySize;
    size_t             keyCapacity;
    unsigned long long hash;
} DataHeader;

typedef struct KstlHashMap {
    size_t              loadFactorSize;
    KstlHashMapSettings settings;
    KstlVector*         buckets;
    KstlList*           recycle;
    size_t              size;
} KstlHashMap;

static KstlVector* InitBuckets(KstlHashMapSettings settings, size_t amount)
{
    KstlList         list;
    KstlListSettings listSettings = {.itemSize = sizeof(DataHeader) + settings.itemSize};
    KstlListInitStruct(&list, listSettings);

    KstlVector* buckets = KstlVectorInitWithCapacityAndValue(sizeof(KstlList), amount, &list);

    return buckets;
}

static void ReleaseDataHeader(DataHeader* header)
{
    KstlRelease(header->key);
}

static int ReleaseBuckets(KstlVector* buckets)
{
    KV_FOR_EACH(buckets, kvIt, kvEnd)
    {
        KstlList* list = KstlVectorIteratorValue(kvIt);
        if (!list || KstlListEmpty(list)) {
            continue;
        }

        KL_FOR_EACH(list, lIt, lEnd)
        {
            DataHeader* header = KstlListIteratorValue(lIt);
            ReleaseDataHeader(header);
        }
        KstlListRelease(list);
    }
    KstlVectorRelease(buckets);

    return 0;
}

static unsigned long long Hash(const void* buf, size_t len)
{
    // FNV-1a Hash algorithm
    unsigned long long hval = 0;
    unsigned char*     bp   = (unsigned char*)buf;
    unsigned char*     be   = bp + len;

    while (bp < be) {
        hval ^= (unsigned long long)*bp++;
        hval += (hval << 1) + (hval << 4) + (hval << 5) + (hval << 7) + (hval << 8) + (hval << 40);
    }

    return hval;
}

static int FillDataHeader(DataHeader* header, const void* key, size_t keySize, unsigned long long hash)
{
    if (header->key && header->keyCapacity < keySize) {
        KstlRelease(header->key);
        header->key = NULL;
    }

    if (!header->key) {
        header->key = KstlAllocate(keySize);
        if (!header->key) {
            return -1;
        }
        header->keyCapacity = keySize;
    }

    memcpy(header->key, key, keySize);

    header->keySize = keySize;
    header->hash    = hash;

    return 0;
}

static KstlList* FindBucket(KstlVector* buckets, unsigned long long hashValue)
{
    size_t bucketIndex = hashValue % KstlVectorSize(buckets);
    return KstlVectorReadAt(buckets, bucketIndex);
}

static KstlVectorIterator FindBucketIt(KstlVector* buckets, unsigned long long hashValue)
{
    size_t bucketIndex = hashValue % KstlVectorSize(buckets);
    return KstlVectorIteratorFromIndex(buckets, bucketIndex);
}

static size_t GetBucketMultiplier(size_t currentSize)
{
    if (currentSize < BUCKETS_MULTIPLIER_GROWING_EDGE) {
        return FAST_GROWING_BUCKETS_MULTIPLIER;
    } else {
        return SLOW_GROWING_BUCKETS_MULTIPLIER;
    }
}

static void UserDealloc(KstlHashMap* handle, KstlListNode* node)
{
    if (handle->settings.dealloc.deallocate) {
        void* context = handle->settings.dealloc.deallocatorContext;
        void* pItem   = (DataHeader*)KstlListPayload(node) + 1;
        handle->settings.dealloc.deallocate(context, pItem);
    }
}

static void Recycle(KstlHashMap* handle, KstlList* bucket, KstlListNode* node)
{
    UserDealloc(handle, node);
    KstlListExtractNode(bucket, node);
    KstlListPushNodeBack(handle->recycle, node);
}

static int Rehash(KstlHashMap* handle)
{
    size_t      currentBucketsAmount = KstlVectorSize(handle->buckets);
    size_t      bucketMultiplier     = GetBucketMultiplier(currentBucketsAmount);
    KstlVector* newBuckets           = InitBuckets(handle->settings, currentBucketsAmount * bucketMultiplier);
    if (!newBuckets) {
        return -1;
    }

    KstlVectorIterator it  = KstlVectorIteratorBegin(handle->buckets);
    KstlVectorIterator end = KstlVectorIteratorEnd(handle->buckets);
    while (!KstlVectorIteratorEqual(it, end)) {
        KstlList* list = (KstlList*)KstlVectorIteratorValue(it);
        while (!KstlListEmpty(list)) {
            DataHeader*   header  = KstlListFront(list);
            KstlList*     newList = FindBucket(newBuckets, header->hash);
            KstlListNode* front   = KstlListPopNodeFront(list);
            KstlListPushNodeBack(newList, front);
        }
        KstlVectorIteratorNext(&it);
    }

    ReleaseBuckets(handle->buckets);
    handle->buckets = newBuckets;

    handle->loadFactorSize = (size_t)(LOAD_FACTOR * KstlVectorCapacity(handle->buckets));
    return 0;
}

static bool CompareKeys(const void* left, size_t leftSize, const void* right, size_t rightSize)
{
    if (leftSize != rightSize) {
        return false;
    }

    return memcmp(left, right, leftSize) == 0;
}

static KstlListIterator FindInBucket(KstlList* bucket, const void* key, size_t keySize)
{
    KstlListIterator it  = KstlListIteratorBegin(bucket);
    KstlListIterator end = KstlListIteratorEnd(bucket);

    while (!KstlListIteratorEqual(it, end)) {
        DataHeader* header = KstlListIteratorValue(it);
        if (CompareKeys(header->key, header->keySize, key, keySize)) {
            return it;
        }

        KstlListIteratorNext(&it);
    }

    return end;
}

static void* FindDataInBucket(KstlList* bucket, const void* key, size_t keySize)
{
    KstlListIterator it = FindInBucket(bucket, key, keySize);
    if (KstlListIteratorEqual(it, KstlListIteratorEnd(bucket))) {
        return NULL;
    }

    return KstlListIteratorValue(it);
}

KstlHashMap* KstlHashMapInit(KstlHashMapSettings settings)
{
    KstlHashMap* handle = KstlAllocate(sizeof(KstlHashMap));
    if (!handle) {
        return NULL;
    }

    memset(handle, 0, sizeof(KstlHashMap));

    handle->loadFactorSize = (size_t)(LOAD_FACTOR * INITIAL_BUCKETS_AMOUNT);
    handle->settings       = settings;

    handle->buckets = InitBuckets(settings, INITIAL_BUCKETS_AMOUNT);
    if (!handle->buckets) {
        KstlHashMapRelease(handle);
        return NULL;
    }

    KstlListSettings recycleSettings = {.itemSize = sizeof(DataHeader) + settings.itemSize};
    handle->recycle                  = KstlListInit(recycleSettings);
    if (!handle->recycle) {
        KstlHashMapRelease(handle);
        return NULL;
    }

    return handle;
}

size_t KstlHashMapSize(KstlHashMap* handle)
{
    return handle->size;
}

size_t KstlHashMapBucketCount(KstlHashMap* handle)
{
    return KstlVectorSize(handle->buckets);
}

void KstlHashMapRelease(KstlHashMap* handle)
{
    if (!handle) {
        return;
    }

    if (handle->buckets) {
        ReleaseBuckets(handle->buckets);
    }

    if (handle->recycle) {
        KL_FOR_EACH(handle->recycle, lIt, lEnd)
        {
            DataHeader* header = KstlListIteratorValue(lIt);
            ReleaseDataHeader(header);
        }
        KstlListRelease(handle->recycle);
    }

    KstlRelease(handle);
}

KstlHashMapIterator KstlHashMapAdd(KstlHashMap* handle, const void* key, const size_t keySize, const void* value)
{
    if (!handle) {
        return KstlHashMapIteratorEnd(handle);
    }

    if (handle->size >= handle->loadFactorSize) {
        Rehash(handle);
    }

    unsigned long long hashValue = Hash(key, keySize);
    KstlVectorIterator bucketIt  = FindBucketIt(handle->buckets, hashValue);
    KstlListIterator   listIt    = FindInBucket(KstlVectorIteratorValue(bucketIt), key, keySize);

    DataHeader*         payload = NULL;
    KstlHashMapIterator it      = {.handle = handle, .bucketIt = bucketIt};
    if (KstlListIteratorEqual(listIt, KstlListIteratorEnd(KstlVectorIteratorValue(bucketIt)))) {
        if (!KstlListEmpty(handle->recycle)) {
            KstlListNode* node = KstlListPopNodeFront(handle->recycle);
            KstlListPushNodeBack(KstlVectorIteratorValue(bucketIt), node);
            payload = KstlListPayload(node);
        } else {
            payload = KstlListPushRawBack(KstlVectorIteratorValue(bucketIt));
        }

        int status = FillDataHeader(payload, key, keySize, hashValue);
        if (status) {
            return KstlHashMapIteratorEnd(handle);
        }

        it.itemIt = KstlListIteratorReverseBegin(KstlVectorIteratorValue(bucketIt));
        handle->size++;
    } else {
        payload   = KstlListIteratorValue(listIt);
        it.itemIt = listIt;
    }

    memcpy(payload + 1, value, handle->settings.itemSize);

    return it;
}

int KstlHashMapGet(KstlHashMap* handle, const void* key, const size_t keySize, void* out)
{
    unsigned long long hashValue = Hash(key, keySize);
    KstlList*          bucket    = FindBucket(handle->buckets, hashValue);

    DataHeader* dataHeader = FindDataInBucket(bucket, key, keySize);
    if (!dataHeader) {
        return -1;
    }

    memcpy(out, dataHeader + 1, handle->settings.itemSize);

    return 0;
}

KstlHashMapIterator KstlHashMapFind(KstlHashMap* handle, const void* key, const size_t keySize)
{
    KstlHashMapIterator it = {.handle = handle};

    unsigned long long hashValue   = Hash(key, keySize);
    size_t             bucketIndex = hashValue % KstlVectorSize(handle->buckets);

    it.bucketIt = KstlVectorIteratorFromIndex(handle->buckets, bucketIndex);
    if (KstlVectorIteratorEqual(it.bucketIt, KstlVectorIteratorEnd(handle->buckets))) {
        return KstlHashMapIteratorEnd(handle);
    }

    KstlList* bucket = KstlVectorIteratorValue(it.bucketIt);
    it.itemIt        = FindInBucket(bucket, key, keySize);
    if (KstlListIteratorEqual(it.itemIt, KstlListIteratorEnd(bucket))) {
        return KstlHashMapIteratorEnd(handle);
    }

    return it;
}

int KstlHashMapErase(KstlHashMap* handle, const void* key, const size_t keySize)
{
    KstlHashMapIterator it = KstlHashMapFind(handle, key, keySize);
    if (KstlHashMapIteratorEqual(it, KstlHashMapIteratorEnd(handle))) {
        return -1;
    }

    KstlHashMapDelete(handle, it);
    return 0;
}

int KstlHashMapDelete(KstlHashMap* handle, KstlHashMapIterator it)
{
    if (handle != it.handle) {
        return -1;
    }

    Recycle(handle, it.bucketIt.pointer, it.itemIt.current);
    handle->size--;
    return 0;
}

void KstlHashMapClear(KstlHashMap* handle)
{
    KV_FOR_EACH(handle->buckets, it, end)
    {
        KstlList* list = KstlVectorIteratorValue(it);
        if (!list) {
            continue;
        }

        while (!KstlListEmpty(list)) {
            Recycle(handle, list, list->head);
        }
    }

    handle->size = 0;
}

KstlHashMapIterator KstlHashMapIteratorBegin(KstlHashMap* handle)
{
    for (KstlVectorIterator it = KstlVectorIteratorBegin(handle->buckets);
         !KstlVectorIteratorEqual(it, KstlVectorIteratorEnd(handle->buckets));
         KstlVectorIteratorNext(&it)) {
        KstlList* list = KstlVectorIteratorValue(it);
        if (!KstlListEmpty(list)) {
            KstlHashMapIterator khmIt = {.handle = handle, .bucketIt = it, .itemIt = KstlListIteratorBegin(list)};
            return khmIt;
        }
    }

    return KstlHashMapIteratorEnd(handle);
}

KstlHashMapIterator KstlHashMapIteratorEnd(KstlHashMap* handle)
{
    KstlHashMapIterator khmIt = {.handle = handle, .bucketIt = KstlVectorIteratorEnd(handle->buckets)};
    khmIt.bucketIt            = KstlVectorIteratorEnd(handle->buckets);
    return khmIt;
}

void KstlHashMapIteratorNext(KstlHashMapIterator* it)
{
    KstlListIteratorNext(&it->itemIt);

    if (!KstlListIteratorEqual(it->itemIt, KstlListIteratorEnd(it->itemIt.handle))) {
        return;
    }

    for (KstlVectorIteratorNext(&it->bucketIt);
         !KstlVectorIteratorEqual(it->bucketIt, KstlVectorIteratorEnd(it->handle->buckets));
         KstlVectorIteratorNext(&it->bucketIt)) {
        KstlList* list = KstlVectorIteratorValue(it->bucketIt);
        if (!KstlListEmpty(list)) {
            it->itemIt = KstlListIteratorBegin(list);
            return;
        }
    }

    it->itemIt.handle = NULL;
}

bool KstlHashMapIteratorEqual(KstlHashMapIterator current, KstlHashMapIterator end)
{
    return current.handle == end.handle && KstlVectorIteratorEqual(current.bucketIt, end.bucketIt) &&
           KstlListIteratorEqual(current.itemIt, end.itemIt);
}

void* KstlHashMapIteratorValue(KstlHashMapIterator current)
{
    DataHeader* header = KstlListIteratorValue(current.itemIt);
    return header + 1;
}

void* KstlHashMapIteratorKey(KstlHashMapIterator current)
{
    DataHeader* header = KstlListIteratorValue(current.itemIt);
    return header->key;
}

size_t KstlHashMapIteratorKeySize(KstlHashMapIterator current)
{
    DataHeader* header = KstlListIteratorValue(current.itemIt);
    return header->keySize;
}

size_t KstlHashMapRecycleCapacity(KstlHashMap* handle)
{
    return KstlListSize(handle->recycle);
}